# %%
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
import tensorflow as tf
from tensorflow import keras
import numpy as np
import getopt, sys
from PIL import Image
tf.config.verbose = 0
physical_devices = tf.config.list_physical_devices('GPU')
tf.config.experimental.set_memory_growth(physical_devices[0], True)
# %%
model = keras.models.load_model('pic2vec.h5')
input_files = []
# %%
try:
  opts, args = getopt.getopt(sys.argv[1:], 'i:')
  for k, v in opts:
    if k == '-i':
      input_files.append(v)
except getopt.GetoptError as err:
  print(err)

if len(input_files) != 2:
  print('need 2 files to compare')
  sys.exit(-1)


# %%
def file_to_data(fpath):
  img = Image.open(fpath)
  x, y = img.size
  r = max(x, y)
  img = img.crop((0, 0, r, r))
  img = img.resize((32, 32))
  return np.asarray(img).reshape(1, 32, 32, 3)


# %%
pics = []
for f in input_files:
  pics.append(file_to_data(f))

ds = np.concatenate(pics)
res = model.predict(ds)

dist = np.linalg.norm(res[0] - res[1])
print(dist)
